package nig.mf.plugin.sink;

import arc.archive.ArchiveInput;
import arc.archive.ArchiveRegistry;
import arc.mf.plugin.DataSinkImpl;
import arc.mf.plugin.PluginLog;
import arc.mime.NamedMimeType;
import arc.streams.LongInputStream;
import arc.streams.StreamCopy;
import arc.xml.XmlDoc.Element;
import ch.ethz.ssh2.Connection;
import ch.ethz.ssh2.SCPClient;
import ch.ethz.ssh2.SCPOutputStream;
import ch.ethz.ssh2.Session;
import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.logging.Level;
import java.util.logging.Logger;
import java.util.Map;
import java.util.concurrent.Callable;
import nig.mf.plugin.sink.SinkLogger.TransferContext;
import nig.ssh.SSH;
import static nig.ssh.Util.illegal;


/**
 * A custom sink that transfers local files to remote SCP destinations.
 * 
 * @author King Lung Chiu
 */
public class SCPSink implements DataSinkImpl
{
   private final SinkLogger                       logger;
   private final Map<TransferContext, Connection> connections;
   public SCPSink() {
      logger      = new SinkLogger(this.getClass().getCanonicalName());
      connections = new ConcurrentHashMap<TransferContext, Connection>();
   }


	@Override public String[]  acceptedTypes() throws Throwable { return null; }
	@Override public int       compressionLevelRequired()       { return -1; }
	@Override public String    description() throws Throwable   { return "SCP Sink"; }
	@Override public String    type() throws Throwable          { return "scp"; }
	@Override public Map       parameterDefinitions() throws Throwable {
		return RemoteSinkParam.definitions();
	}

   @Override
   public Object beginMultiple(Map<String, String> params) throws Throwable {
      if(params == null) illegal("params should not be null!");
      
      TransferContext ctx = logger.beginMultiple(params);
      connections.put(ctx, createConnection(params));
      return ctx;
   }

   @Override
   public void endMultiple(Object ctxKey) throws Throwable {
      if(ctxKey == null)
         illegal("ctxkey cannot be null!");
      if(!logger.sessionExists(ctxKey))
         illegal("ctxKey does not exist in internal logger's sesssions list! Has it been previously removed?");
      
      TransferContext session = logger.getSession(ctxKey);
      Connection con = connections.remove(session);
      if(con != null) {
         con.close();         
      } else {
         logger.log(
            PluginLog.ERROR,
            "endMultiple(ctxKey):\n" +
            " connections.remove(session) unexpectedly returned null!"
         );         
      }
      logger.endMultiple(ctxKey);
   }

	@Override
	public void shutdown() throws Throwable {
      logger.shutdown(new Callable<String>() {
         @Override public String call() throws Exception {
            StringBuilder msg = new StringBuilder()
               .append("   ")
               .append(this.getClass().getSimpleName())
               .append(" ...");
            if(connections.isEmpty())
               return msg.append("\n   no open connections to close").toString();

            for(Connection con: connections.values()) {
               try {
                  con.close();
                  msg.append("\n   connection closed for: ").append(con.getHostname());
               } catch(Exception error) {
                  msg.append("\n   connection close failed for: ").append(con.getHostname());
                  logger.log(
                     PluginLog.ERROR,
                     "couldn't close connection for '" + con.getHostname() + "'",
                     error
                  );
               }
            }
            connections.clear();
            return msg.toString();
         }
      });
	}

	@Override
	public void consume(
      Object multiCtx,
      String path,
      Map<String, String> params,
      Element userMeta,
      Element meta,
      LongInputStream in,
      String appMimeType,
      String streamMimeType,
      long length
   ) throws Throwable {
      logger.consume(multiCtx, path, params, userMeta, meta, in, appMimeType, streamMimeType, length);
      if(path == null || path.isEmpty())
         illegal("Missing argument: path");
      if(!RemoteSinkParam.path.existsIn(params))
         illegal("Missing sink params: path MUST BE SUPPLIED");
      if(in == null)
         illegal("Missing argument: in");
      if(length < 0)
         illegal("Input stream's length cannot be negative.");

      // if called, createConnection() also checks whether host,
      // hostFingerprint, username and privateKey are supplied
      final Connection con = logger.sessionExists(multiCtx)?
         connections.get(logger.getSession(multiCtx)):
         createConnection(params);
      Session ssh = null;
      SCPClient scp = null;
      
      try {
         ssh = con.openSession();
         scp = con.createSCPClient();
//         final byte[] buff = new byte[RemoteSinkParam.buffSizeFrom(params)];
         final String baseDir = RemoteSinkParam.path.from(params);
         final String fullPath =
            (baseDir + "/" + path).replace("\\\\", "/").replace("\\", "/");
         
         if((streamMimeType != null) && ArchiveRegistry.isAnArchive(streamMimeType)) {
            logger.log(
               PluginLog.INFORMATION,
               "consume() treating input stream as an archive"
            );
            writeArchive(
               ArchiveRegistry.createInput(in, new NamedMimeType(streamMimeType)),
               con,
               (fullPath + "/").replaceAll("/+", "/")
            );
         } else {
            String dir = "";        // fullPath may just be a file, hence empty dir
            String name = fullPath; // fullPath may just be a file, hence = name
            int pos = fullPath.lastIndexOf("/");
            if(pos > -1) {
               dir = fullPath.substring(0, pos+1);
               name = fullPath.substring(pos+1);
            }
            logger.log(
               PluginLog.INFORMATION,
               "consume() treating input stream as a raw binary stream\n" +
               "   dir : " + dir +
               "   name: " + name +
               "\n"
            );
//            SSH.mkdirs(ssh, dir);
            writeFile(con, in, length, name, dir);
         }
      } finally {
         // always close the SSH session
         if(ssh != null) ssh.close();

         // close call-specific connection
         if(con != null && !logger.sessionExists(multiCtx)) {
            con.close();
         }         
      }
   }

   private void writeArchive(
      final ArchiveInput archive,
      final Connection con,
      final String baseDir   // TODO: make sure ends with '/' and no duplicate /'s
//      final byte[] buff    // currently unused
   ) throws IOException, Throwable {
      try {
         ArchiveInput.Entry entry = archive.next();
         while(entry != null) {
            // !!!! redundant if file name already contains path info !!!!
            if(entry.isDirectory()) {
               SSH.mkdirs(con, fullPath(baseDir, entry.name()));
            }
            else {
               String name = entry.name();
               String dir = baseDir;
               int pos = name.lastIndexOf("/");
               if(pos > -1) {
                  dir = fullPath(baseDir, name.substring(0, pos + 1));
                  name = name.substring(pos + 1);
               }
               writeFile(con, entry.stream(), entry.size(), name, dir);
            }
            entry = archive.next();
         }

      } catch (Throwable ex) {
         Logger.getLogger(this.getClass().getName()).log(Level.SEVERE, null, ex);
         throw ex;
      }
   }

   private void writeFile(
      final Connection con,
      final LongInputStream lis,
      final long lisSize,
      final String name,
      final String dir
   ) throws Throwable, IOException {
      
      SCPOutputStream os = null;
      
      try {
         // ensure all parent directories are created first
         SSH.mkdirs(con, dir);
         
         // then create and write to the rmote file
         SCPClient scp = con.createSCPClient();
         os = scp.put(name, lisSize, dir, "0600");
         // use this if don't need to adjust buffer size
         StreamCopy.copy(lis, os);

      } finally {
         if(os != null)
            os.close();
      }      
   }
//   private void writeFile(
//      final SCPClient scp,
//      final LongInputStream lis,
//      final long lisSize,
//      final String name,
//      final String dir
////      final byte[] buff  // currently unused
//   ) throws Throwable, IOException {
//      SCPOutputStream os = scp.put(name, lisSize, dir, "0600");
//      try {
//          // use this if don't need to adjust buffer size
//          StreamCopy.copy(lis, os);
//
////         long transferred = 0;
////         while(transferred < lisSize) {
////            int size = lis.read(buff);
////            if(size > 0) {
////               os.write(buff, 0, size);
////               transferred += size;
////            }
////         }
////         while(lis.hasData()) {
////            int size = lis.read(buff);
////            if(size > 0)
////               os.write(buff, 0, size);
////         }
//      } finally {
//         os.close();
//      }      
//   }

   /**
    * Construct a full path from a base dir and suffix path
    *
    * @param baseDir MUST already contain trailing /
    * @param suffix
    * @return the constructed path
    */
   private String fullPath(String baseDir, String suffix) {
      return new StringBuilder().append(baseDir).append(suffix).toString();
   }

   private Connection createConnection(Map<String, String> params) throws IOException {
      return SSH.authenticatedConnection(
         RemoteSinkParam.host.from(params),
         RemoteSinkParam.portFrom(params),
         RemoteSinkParam.hostFingerprint.from(params),
         RemoteSinkParam.username.from(params),
         RemoteSinkParam.privateKey.from(params).toCharArray()
      );
   }
}